import os
import tqdm

root = "/data/ct/gunman/datasets/COD10k"
train_folder = 'TrainDataset'
output_file = '/data/ct/gunman/UGTR/train_test_file/train_1.lst'

with open(output_file, 'w') as f:
    for Imgfilename in tqdm.tqdm(os.listdir(f"{os.path.join(root, train_folder)}/Imgs")):
        Imgfilename = f"{train_folder}/Imgs/{Imgfilename}"
        GTfilename = f"{train_folder}/GT/{Imgfilename.split('/')[-1].split('.')[0]}.png"
        Edgefilename = f"{train_folder}/Edge/{Imgfilename.split('/')[-1].split('.')[0]}.png"
        # 判断GT和Edge是否存在
        if os.path.exists(f"{os.path.join(root, GTfilename)}") and os.path.exists(f"{os.path.join(root, Edgefilename)}"):
            f.write(f"{Imgfilename} {GTfilename} {Edgefilename}\n")